﻿//-----------------------------------------------------------------------
// <copyright file="ShardingConsumerControllerImpl.cs" company="Akka.NET Project">
//     Copyright (C) 2009-2022 Lightbend Inc. <http://www.lightbend.com>
//     Copyright (C) 2013-2025 .NET Foundation <https://github.com/akkadotnet/akka.net>
// </copyright>
//-----------------------------------------------------------------------

#nullable enable
using System;
using System.Collections.Immutable;
using System.Linq;
using Akka.Actor;
using Akka.Delivery;
using Akka.Event;
using Akka.Util.Extensions;

namespace Akka.Cluster.Sharding.Delivery.Internal;

/// <summary>
/// INTERNAL API
/// </summary>
/// <typeparam name="T">The types of messages handled by the ConsumerController</typeparam>
internal class ShardingConsumerController<T> : ReceiveActor, IWithStash, IWithTimers
{
    private const string ShutdownTimeoutTimerKey = nameof(ShutdownTimeoutTimerKey);
    
    private sealed class ShutdownTimeout
    {
        public static readonly ShutdownTimeout Instance = new ();
        private ShutdownTimeout() { }
    }
    
    public ShardingConsumerController(Func<IActorRef, Props> consumerProps,
        ShardingConsumerController.Settings settings)
    {
        ConsumerProps = consumerProps;
        Settings = settings;
        WaitForStart();
    }

    public Func<IActorRef, Props> ConsumerProps { get; }
    public ShardingConsumerController.Settings Settings { get; }

    private readonly ILoggingAdapter _log = Context.GetLogger();

    // gets started asynchronously in the PreStart method
    private IActorRef _consumer = ActorRefs.Nobody;

    /// <summary>
    /// Map of producerControllers to producerIds
    /// </summary>
    public ImmutableDictionary<IActorRef, string> ProducerControllers { get; private set; } =
        ImmutableDictionary<IActorRef, string>.Empty;

    /// <summary>
    /// Map of producerIds to consumerControllers
    /// </summary>
    public ImmutableDictionary<string, IActorRef> ConsumerControllers { get; private set; } =
        ImmutableDictionary<string, IActorRef>.Empty;

    private void WaitForStart()
    {
        Receive<ConsumerController.Start<T>>(start =>
        {
            ConsumerController.AssertLocalConsumer(start.DeliverTo);
            Context.Unwatch(_consumer);
            _consumer = start.DeliverTo;
            Context.Watch(start.DeliverTo);
            Become(Active);
            Stash.UnstashAll();
        });

        Receive<ConsumerController.IConsumerCommand<T>>(_ => { Stash.Stash(); });

        Receive<Terminated>(_ =>
        {
            _log.Debug("Consumer terminated before initialized.");
            Context.Stop(Self);
        });

        Receive<Passivate>(_ => Sender.Equals(_consumer), p =>
        {
            Context.Parent.Tell(p);
        });
        
        ReceiveAny(msg =>
        {
            if (Settings.AllowBypass)
            {
                _consumer.Forward(msg);
            }
            else
            {
                _log.Warning($"Message unhandled [{msg}]. If you need to pass this message to the consumer sharding entity actor, set \"akka.reliable-delivery.sharding.consumer-controller.allow-bypass\" to true");
                Unhandled(msg);
            }
        });
    }

    private void Active()
    {
        Receive<ConsumerController.SequencedMessage<T>>(seqMsg =>
        {
            if (ConsumerControllers.TryGetValue(seqMsg.ProducerId, out var consumerController))
            {
                consumerController.Tell(seqMsg);
                ProducerControllers = UpdatedProducerControllers(seqMsg.ProducerController, seqMsg.ProducerId);
            }
            else
            {
                _log.Debug("Starting ConsumerController for producerId {0}", seqMsg.ProducerId);
                var cc = Context.ActorOf(
                    ConsumerController.Create<T>(Context, seqMsg.ProducerController.AsOption(),
                        Settings.ConsumerControllerSettings), Uri.EscapeDataString($"consumerController-{seqMsg.ProducerId}"));
                Context.Watch(cc);
                cc.Tell(new ConsumerController.Start<T>(_consumer));
                cc.Tell(seqMsg);
                ConsumerControllers = ConsumerControllers.Add(seqMsg.ProducerId, cc);
                ProducerControllers = UpdatedProducerControllers(seqMsg.ProducerController, seqMsg.ProducerId);
            }
        });

        Receive<Terminated>(t => t.ActorRef.Equals(_consumer), _ =>
        {
            _log.Debug("Consumer terminated.");
            
            // Short-circuit shutdown process, just shut down immediately if there's nothing to clean.
            if (ProducerControllers.Count == 0 && ConsumerControllers.Count == 0)
            {
                _log.Debug("ShardingConsumerController terminated.");
                Context.Stop(Self);
            }
            else
            {
                Become(ShuttingDown());
            }
        });

        Receive<Terminated>(t =>
        {
            if (ProducerControllers.TryGetValue(t.ActorRef, out var producer))
            {
                _log.Debug("ProducerController for producerId [{0}] terminated.", producer);
                ProducerControllers = ProducerControllers.Remove(t.ActorRef);
                if (ConsumerControllers.TryGetValue(producer, out var consumerController))
                {
                    consumerController.Tell(ConsumerController.DeliverThenStop<T>.Instance);
                }
            }
            else
            {
                var found = false;
                foreach (var kv in ConsumerControllers.Where(c =>
                             c.Value.Equals(t.ActorRef)))
                {
                    found = true;
                    _log.Debug("ConsumerController for producerId [{0}] terminated.", kv.Key);
                    ConsumerControllers = ConsumerControllers.Remove(kv.Key);
                }

                if (!found)
                {
                    _log.Debug("Unknown [{0}] terminated.", t.ActorRef);
                }
            }
        });

        Receive<Passivate>(_ => Sender.Equals(_consumer), p =>
        {
            Context.Parent.Tell(p);
        });
        
        ReceiveAny(msg =>
        {
            if (Settings.AllowBypass)
            {
                _consumer.Forward(msg);
            }
            else
            {
                _log.Warning($"Message unhandled [{msg}]. If you need to pass this message to the consumer sharding entity actor, set \"akka.reliable-delivery.sharding.consumer-controller.allow-bypass\" to true");
                Unhandled(msg);
            }
        });
    }

    // Shutdown state after `_consumer` actor is downed.
    private Action ShuttingDown()
    {
        // start a 3-seconds shutdown timeout timer
        Timers.StartSingleTimer(ShutdownTimeoutTimerKey, ShutdownTimeout.Instance, TimeSpan.FromSeconds(3), Self);

        _log.Debug("Shutting down child controllers");

        foreach (var p in ProducerControllers.Keys)
            Context.Unwatch(p);
        ProducerControllers = ImmutableDictionary<IActorRef, string>.Empty;
        
        foreach (var c in ConsumerControllers.Values.Distinct())
            Context.Stop(c);
        
        return () =>
        {
            Receive<ConsumerController.SequencedMessage<T>>(seqMsg =>
            {
                var messageType = seqMsg.Message.Chunk.HasValue 
                    ? $"Manifest: {seqMsg.Message.Chunk.Value.Manifest}, SerializerId: {seqMsg.Message.Chunk.Value.SerializerId}"  
                    : seqMsg.Message.Message?.GetType().FullName ?? "Unknown type";
                _log.Warning("Message [{0}] from [{1}] is being ignored because ShardingConsumerController is shutting down.", messageType, seqMsg.ProducerId);
            });

            Receive<ShutdownTimeout>(_ =>
            {
                // We somehow could not terminate cleanly within 3 seconds, shutdown immediately
                _log.Warning("ShardingConsumerController cleanup timed out, force terminating.");
                Context.Stop(Self);
            });

            Receive<Terminated>(t =>
            {
                var removeList = ConsumerControllers
                    .Where(kv => kv.Value.Equals(t.ActorRef))
                    .Select(kv => kv.Key)
                    .ToArray();
                    
                if(removeList.Length > 0)
                {
                    foreach (var key in removeList)
                        _log.Debug("ConsumerController for producerId [{0}] terminated.", key);
                        
                    ConsumerControllers = ConsumerControllers.RemoveRange(removeList);
                }
                
                if (ProducerControllers.Count > 0 || ConsumerControllers.Count > 0)
                    return;

                _log.Debug("ShardingConsumerController terminated.");
                Context.Stop(Self);
            });
        };
    }

    private ImmutableDictionary<IActorRef, string> UpdatedProducerControllers(IActorRef producerController,
        string producer)
    {
        if (ProducerControllers.ContainsKey(producerController))
            return ProducerControllers;
        Context.Watch(producerController);
        return ProducerControllers.Add(producerController, producer);
    }

    protected override void PreStart()
    {
        var self = Self;
        _consumer = Context.ActorOf(ConsumerProps(self), "consumer");
        Context.Watch(_consumer);
    }

    public IStash Stash { get; set; } = null!;
    public ITimerScheduler Timers { get; set; } = null!;
}
